/*=====================================================================*
 *                   Copyright (C) 2011 Paul Mineiro                   *
 * All rights reserved.                                                *
 *                                                                     *
 * Redistribution and use in source and binary forms, with             *
 * or without modification, are permitted provided that the            *
 * following conditions are met:                                       *
 *                                                                     *
 *     * Redistributions of source code must retain the                *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer.                                       *
 *                                                                     *
 *     * Redistributions in binary form must reproduce the             *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer in the documentation and/or            *
 *     other materials provided with the distribution.                 *
 *                                                                     *
 *     * Neither the name of Paul Mineiro nor the names                *
 *     of other contributors may be used to endorse or promote         *
 *     products derived from this software without specific            *
 *     prior written permission.                                       *
 *                                                                     *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND              *
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,         *
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES               *
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE             *
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER               *
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,                 *
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES            *
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE           *
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR                *
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF          *
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT           *
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY              *
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE             *
 * POSSIBILITY OF SUCH DAMAGE.                                         *
 *                                                                     *
 * Contact: Paul Mineiro <paul@mineiro.com>                            *
 *=====================================================================*/

#include "mstep.h"
#include "qfunc.h"

#include <math.h>
#include <stdlib.h>
#include <string.h>

static void
make_index (unsigned int* index,
            unsigned int  n_ratings)
{
  index[0] = 0;

  for (unsigned int i = 1; i < n_ratings; ++i)
    {
      unsigned int j = (unsigned int) ((i + 1) * drand48 ());
      index[i] = index[j];
      index[j] = i;
    }
}

void
mstep_sgd (float*                               loga,
           unsigned int                         n_workers,
           unsigned int                         rank,
           float*                               gamma,
           float*                               logbeta,
           unsigned int                         n_images,
           float*                               c,
           unsigned int                         n_labels,
           float*                               d,
           const float*                         pz,
           Rating*                              ratings,
           unsigned int                         n_ratings,
           const unsigned int*                  image_ratings,
           const unsigned int*                  worker_ratings,
           const ParameterizedDistribution*     priorloga,
           const Distribution*                  priorgamma,
           const NormalDistribution*            priorlogbeta,
           const Distribution*                  priorc,
           const Distribution*                  priord,
           unsigned int                         n_passes,
           float                                eta)   /* e.g, eta = 1 */
{
  float dloga[n_workers * rank];
  float dgamma[rank];
  float dlogbeta[n_images];
  float dc[n_labels * rank];
  float dd[n_labels * rank];
  unsigned int index[n_ratings];

  for (unsigned int k = 0; k < n_passes; ++k)
    {
      make_index (index, n_ratings);

      for (unsigned int n = 0; n < n_ratings; ++n)
        {
          float thiseta = eta / sqrt (1.0 + n + k * n_ratings);
          unsigned int image = ratings[index[n]].image;
          unsigned int worker = ratings[index[n]].worker;

          memset (dloga + worker * rank, 0, rank * sizeof (dloga[0]));
          memset (dgamma, 0, rank * sizeof (dgamma[0]));
          dlogbeta[image] = 0;
          memset (dc, 0, rank * n_labels * sizeof (dc[0]));
          memset (dd, 0, rank * n_labels * sizeof (dd[0]));

          dqfunc_single (index[n], image_ratings[image], worker_ratings[worker],
                         dloga, dgamma, dlogbeta, dc, dd, loga, 
                         rank, gamma, logbeta, c, n_labels, d, pz, 
                         ratings, n_ratings, priorloga, priorgamma, 
                         &priorlogbeta->base, priorc, priord);

          for (unsigned int i = 0; i < rank; ++i)
            {
              loga[worker * rank + i] += thiseta * dloga[worker * rank + i];
            }

          for (unsigned int i = 0; i < rank; ++i)
            {
              gamma[i] += thiseta * dgamma[i];
            }

          logbeta[image] += thiseta * dlogbeta[image];

          for (unsigned int i = 0; i < n_labels * rank; ++i)
            {
              c[i] += thiseta * dc[i];
              d[i] += thiseta * dd[i];
            }
        }
    }
}
